import torch

torch.device("cuda")

x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

# 定义参数 w，初值 1.0，形状是标量张量（实际上是 shape=(1,)）。
# requires_grad=True 表示 跟踪梯度，后向传播时会计算 ∂loss/∂w 并存到 w.grad。
w = torch.Tensor([1.0])
w.requires_grad = True


# 线性模型的前向计算：ŷ = x * w。
# 注意：你在循环里传进来的 x 是 Python 浮点数，不是 Tensor。PyTorch 会在运算时把它当作常数参与广播，与 w（Tensor）相乘，结果仍是 Tensor。
# 因为只有 w 需要梯度，所以 autograd 只会对 w 建图、求导。
def forward(x):
    return x * w


# 损失定义为平方误差：(ŷ - y)^2。
# 这里 y 也是 Python 浮点数，会被当常数参与计算。
# 返回的是一个 Tensor（标量），这对 l.backward() 很重要：标量张量才能直接调用 backward() 不带参数。
def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) ** 2


print("predict (before training)", 4, forward(4))

for epoch in range(100):
    for x, y in zip(x_data, y_data):
        # 基于当前 w 计算该样本的损失 l（标量 Tensor）。
        # 此时 autograd 已记录了从 w 到 l 的计算图。
        l = loss(x, y)
        # 反向传播：计算 ∂l/∂w 并累计到 w.grad。
        l.backward()
        print("\tgrad: ", x, y, w.grad.item())
        # SGD 参数更新：w ← w − η * grad，学习率 η=0.01。
        w.data = w.data - 0.01 * w.grad.data
        # 梯度清零：因为 PyTorch 的梯度默认会累加（accumulate），每个 step 后要把 w.grad 清零，避免下一步把旧梯度叠加进去。
        w.grad.data.zero_()
    print("progress:", epoch, l.item())

print("predict (after training)", 4, forward(4))

# requires_grad=True：告诉 autograd 需要对这个张量求导。这里只对 w 求导。
# 计算图：forward → loss 会构建从 w 到 l 的图，l.backward() 沿图反传，把梯度写到 w.grad。
# 梯度累加：每次 backward() 都会把新梯度加到旧的 w.grad 上，所以每步更新后都需要 grad.zero_()。
# 更新与 no_grad：更新参数时应该关闭梯度跟踪，否则会把“更新操作”也记进图里（低概率引发问题但不是好习惯）。
